Repository navigation
refactor: defer nufftax/jax and numba imports to first use - #451
Merged
Merged
Conversation
…#1505) _load_nufftax() anchors the nufftax import + 0.6.x batcher shim at TransformerNUFFT use (init + the three module-global-dereferencing methods, covering unpickled worker instances); numba_util.jit lazily materializes all decorated functions on first call, rebinding module globals so nopython cross-calls resolve. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Defers PyAutoArray's two heavy import-time costs to first use, as part of the cross-repo import-time task (PyAutoFit#1505): the nufftax import (which pulls jax, ~0.7s) and numba jit decoration (~0.2s).
import autoarraydrops to ~1.15s with neither jax, nufftax nor numba loading on bare import.operators/transformer.py: the module-levelimport nufftaxand the version-gated 0.6.x batcher shim move into an idempotent_load_nufftax(), called fromTransformerNUFFT.__init__and the three methods that dereference the module global (_forward_native,image_from,transform_mapping_matrix) — instances unpickled in multiprocessing workers never re-run__init__, so the method anchors are load-bearing.numba_util.py:jit()queues functions and hands them to numba on the first call of any decorated function, rebinding each defining module's global to the real dispatcher so nopython cross-calls (5 exist) still resolve at compile time.ModuleNotFoundError → plain functionfallback preserved; config-driven nopython/cache/parallel options unchanged.API Changes
None — internal changes only. All 29
@numba_util.jit()functions andTransformerNUFFTbehave identically; laziness only moves when the underlying imports happen.See full details below.
Test Plan
TransformerNUFFTnumpy/jax visibilities agree exactly;transform_mapping_matrix+image_fromverified; pickled-instance-with-reset-module-state (worker simulation) worksbinned_mask_from→binned_image_from) compiles and runs; module global rebinds toCPUDispatcherFull API Changes (for automation & release notes)
Changed Behaviour
autoarray.numba_util.jit— decorated functions are lightweight wrappers until first call (then real numba dispatchers, module globals rebound). Dispatcher-only attributes (e.g..py_func) are unavailable before first call; no users exist in the organism.autoarray.operators.transformer—nufftax(and jax via it) import on firstTransformerNUFFTuse instead of at module import; a missing nufftax now errors at construction (samenufftax_exception()), not silently at import.Migration
Generated by the PyAutoLabs agent workflow.